from PIL import Image
import os
from tqdm import tqdm
import numpy as np

source_dir = "gaofen"

target_dir = "gaofen_aug"



def aug(name):
    img_1 = Image.open(os.path.join(source_dir, "image", f"{name}_1.png"))
    img_2 = Image.open(os.path.join(source_dir, "image", f"{name}_2.png"))
    img_1_label = Image.open(os.path.join(source_dir, "gt", f"{name}_1_label.png"))
    img_2_label = Image.open(os.path.join(source_dir, "gt", f"{name}_2_label.png"))
    label = Image.open(os.path.join(source_dir, "gt", f"{name}_change.png"))
    
    # origin
    img_1.save(os.path.join(target_dir, "image", f"{name}_1.png"))
    img_2.save(os.path.join(target_dir, "image", f"{name}_2.png"))
    img_1_label.save(os.path.join(target_dir, "gt", f"{name}_1_label.png"))
    img_2_label.save(os.path.join(target_dir, "gt", f"{name}_2_label.png"))
    label.save(os.path.join(target_dir, "gt", f"{name}_change.png"))

    # rot 90
    img_1.rotate(90).save(os.path.join(target_dir, "image", f"{name}_90_1.png"))
    img_2.rotate(90).save(os.path.join(target_dir, "image", f"{name}_90_2.png"))
    img_1_label.rotate(90).save(os.path.join(target_dir, "gt", f"{name}_90_1_label.png"))
    img_2_label.rotate(90).save(os.path.join(target_dir, "gt", f"{name}_90_2_label.png"))
    label.rotate(90).save(os.path.join(target_dir, "gt", f"{name}_90_change.png"))
    
    
    # rot 180
    img_1.rotate(180).save(os.path.join(target_dir, "image", f"{name}_180_1.png"))
    img_2.rotate(180).save(os.path.join(target_dir, "image", f"{name}_180_2.png"))
    img_1_label.rotate(180).save(os.path.join(target_dir, "gt", f"{name}_180_1_label.png"))
    img_2_label.rotate(180).save(os.path.join(target_dir, "gt", f"{name}_180_2_label.png"))
    label.rotate(180).save(os.path.join(target_dir, "gt", f"{name}_180_change.png"))
    
    # rot 270
    img_1.rotate(270).save(os.path.join(target_dir, "image", f"{name}_270_1.png"))
    img_2.rotate(270).save(os.path.join(target_dir, "image", f"{name}_270_2.png"))
    img_1_label.rotate(270).save(os.path.join(target_dir, "gt", f"{name}_270_1_label.png"))
    img_2_label.rotate(270).save(os.path.join(target_dir, "gt", f"{name}_270_2_label.png"))
    label.rotate(270).save(os.path.join(target_dir, "gt", f"{name}_270_change.png"))
    
    # vertical
    img_1.transpose(Image.FLIP_TOP_BOTTOM).save(os.path.join(target_dir, "image", f"{name}_v_1.png"))
    img_2.transpose(Image.FLIP_TOP_BOTTOM).save(os.path.join(target_dir, "image", f"{name}_v_2.png"))
    img_1_label.transpose(Image.FLIP_TOP_BOTTOM).save(os.path.join(target_dir, "gt", f"{name}_v_1_label.png"))
    img_2_label.transpose(Image.FLIP_TOP_BOTTOM).save(os.path.join(target_dir, "gt", f"{name}_v_2_label.png"))
    label.transpose(Image.FLIP_TOP_BOTTOM).save(os.path.join(target_dir, "gt", f"{name}_v_change.png"))
    
    # horizontal
    img_1.transpose(Image.FLIP_LEFT_RIGHT).save(os.path.join(target_dir, "image", f"{name}_h_1.png"))
    img_2.transpose(Image.FLIP_LEFT_RIGHT).save(os.path.join(target_dir, "image", f"{name}_h_2.png"))
    img_1_label.transpose(Image.FLIP_LEFT_RIGHT).save(os.path.join(target_dir, "gt", f"{name}_h_1_label.png"))
    img_2_label.transpose(Image.FLIP_LEFT_RIGHT).save(os.path.join(target_dir, "gt", f"{name}_h_2_label.png"))
    label.transpose(Image.FLIP_LEFT_RIGHT).save(os.path.join(target_dir, "gt", f"{name}_h_change.png"))
    
if __name__ == "__main__":
    filenames = set()
    for i in os.listdir(f"{source_dir}/image"):
        if i.endswith("_1.png"):
            filenames.add(i.replace("_1.png", ""))
    filenames = list(filenames)
    for filename in tqdm(filenames):
        aug(filename)
    
    
